from manimlib import *
from collections import deque


class ClickSquare(Button):

    def __init__(self, side_length, char="0", scene=None, **kwargs):
        assert(scene)
        mobject = VGroup(
            Square(side_length, **kwargs),
            Text(char)
        )
        self.char = char
        self.scene_ctx = scene
        super().__init__(mobject, self.on_click, **kwargs)
        self.data["can_move"] = ORIGIN

    def on_click(self, mob):
        if (self.data["can_move"] == ORIGIN).all():
            return
        scene = self.scene_ctx
        distance = scene.square_side_length + scene.square_buff
        scene.play(mob.animate.shift(distance * self.data["can_move"]),
                   run_time=0.4, rate_func=rush_from)
        self.swap_array_item(scene.array, self.char, "x")
        scene.update_can_move()
        # print(scene.array)

    def swap_array_item(self, array: list, i, j):
        a, b = array.index(i), array.index(j)
        array[a], array[b] = array[b], array[a]


class EightPuzzleScene(Scene):

    CONFIG = {
        "square_side_length": 2,
        "square_buff": 0.2,
        "fill_opacity": 0,
        "box_center": LEFT * 2.5,
        "array": ["1", "2", "3", "4", "5", "6", "7", "8", "x"],
        "color_list": [RED, ORANGE, GOLD, YELLOW, GREEN, TEAL, BLUE, PURPLE],
        "on_anim": False
    }

    # def __init__(self, **kwargs):
    #     super().__init__(**kwargs)
    #     digest_config(self, kwargs)

    def init_game_box(self):
        box = Square(3 * self.square_side_length + 4 * self.square_buff)
        box.shift(self.box_center)
        self.add(box)
        self.shift_distance = self.square_side_length + self.square_buff

    def init_squares(self):
        self.squares = Group()
        for i in range(self.array.__len__() - 1):
            self.squares.add(ClickSquare(
                self.square_side_length,
                char=self.array[i],
                scene=self,
                color=self.color_list[i],
                fill_opacity=self.fill_opacity)
            )
        self.squares.arrange_in_grid(3, 3, self.square_buff)
        self.squares.move_to(self.box_center)
        self.add(self.squares)
        # self.update_can_move()     # this line may cause a bug

    def update_can_move(self):
        x_index = self.array.index('x')
        # row, col = x_index // 3, x_index % 3
        for i in range(self.array.__len__()):
            if x_index == i:
                continue
            char = self.array[i]
            if i == x_index - 1:
                self.squares[int(char) - 1].data["can_move"] = RIGHT
            elif i == x_index + 1:
                self.squares[int(char) - 1].data["can_move"] = LEFT
            elif i == x_index - 3:
                self.squares[int(char) - 1].data["can_move"] = DOWN
            elif i == x_index + 3:
                self.squares[int(char) - 1].data["can_move"] = UP
            else:
                self.squares[int(char) - 1].data["can_move"] = ORIGIN

    def init_buttons(self):
        self.info = Text("单击正方形 将它挪到空格处", font="Source Han Serif SC").add_background_rectangle(
            color=BLACK, buff=0.25)
        self.info.to_edge(DOWN)

        play_rec = Text("Play").add_background_rectangle(
            color=BLACK, buff=0.25)
        play_btn = Button(play_rec, self.start_to_play)
        info_rec = Text("Help").add_background_rectangle(
            color=BLACK, buff=0.25)
        info_btn = Button(info_rec, self.show_info)
        solve_rec = Text("Solve").add_background_rectangle(
            color=BLACK, buff=0.25)
        solve_btn = Button(solve_rec, self.solve_problem)
        play_btn.shift(UP*1.2+RIGHT*5)
        info_btn.shift(DOWN*1.2+RIGHT*5)
        solve_btn.shift(RIGHT*5)
        self.add(play_btn, info_btn, solve_btn)

    def start_to_play(self, mob):
        if self.on_anim == True:
            return
        self.on_anim = True
        random.seed(time.time())
        random.shuffle(self.array)
        while not self.sln.has_answer(self.array):
            random.shuffle(self.array)

        pos = dict()
        relative_center = self.box_center + UL * self.shift_distance
        for i in range(len(self.array)):    # i is the index of array
            if self.array[i] == 'x':
                continue
            index = int(self.array[i]) - 1  # the index of square
            row, col = i // 3, i % 3
            pos[index] = relative_center + row * DOWN * \
                self.shift_distance + col * RIGHT * self.shift_distance
        self.play(LaggedStart(*[
            ApplyMethod(self.squares[i].move_to, pos[i],
                        path_arc=PI * (random.random() - 0.5))
            for i in pos
        ]), run_time=1.5, lag_ratio=0.1)
        self.update_can_move()
        self.on_anim = False
        # print(pos)

    def solve_problem(self, mob):
        if self.on_anim == True:
            return
        self.on_anim = True
        print("Solving...")
        ima = time.time()
        steps = self.sln.get_answer("".join(self.array))
        print(steps)
        if steps is None:
            self.on_anim = False
            return
        for i in steps:
            index = int(self.array[i]) - 1
            self.squares[index].on_click(self.squares[index])
        print("Finish!")
        self.on_anim = False

    def show_info(self, mob):
        if self.on_anim == True:
            return
        self.on_anim = True
        self.play(FadeIn(self.info, UP))
        self.wait()
        self.play(FadeOut(self.info, DOWN))
        self.on_anim = False

    def setup(self):
        self.sln = Solution()
        self.init_game_box()
        self.init_squares()
        self.init_buttons()


class Solution:
    """
    传入一个字符串, 返回一个操作序列, 
    这个序列是绝对坐标的点击序列
    """

    def __init__(self):
        self.answer = "12345678x"

    def get_array_that_ai_larger_than_previous_numbers(self, array: list) -> list:
        res = [0]
        for i in range(1, len(array)):
            b_i = 0
            for j in range(i):
                try:
                    if int(array[i]) > int(array[j]):
                        b_i += 1
                except:
                    pass
            res.append(b_i)

        return res

    # 是不是需要忽略 x?
    def has_answer(self, start: str) -> bool:
        sum_array = sum(self.get_array_that_ai_larger_than_previous_numbers(list(start)))
        return sum_array % 2 == 0

    def get_answer(self, start: str) -> list:
        d = { start: [] }
        q = deque()
        q.append(start)
        dx = [0, 0, 1, -1]
        dy = [1, -1, 0, 0]
        result = list()
        while q:
            t = q.popleft()
            distance = d[t].copy()
            if t == self.answer:
                return d[t]
            k = t.index('x')
            x, y = k // 3, k % 3

            tl = list(t)
            for i in range(4):
                a = x + dx[i]
                b = y + dy[i]
                if 0 <= a < 3 and 0 <= b < 3:
                    nk = a * 3 + b
                    tl[nk], tl[k] = tl[k], tl[nk]
                    t = "".join(tl)
                    if t not in d:
                        q.append(t)
                        d[t] = distance + [nk]
                    tl[nk], tl[k] = tl[k], tl[nk]
        return None
